import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.ticker import MaxNLocator
from matplotlib.colors import LinearSegmentedColormap

def plotIntensity(ax, XTicks, YTicks, ResultList):
    #set colormap
    color1 = ["#ffffff","#ff0000","#7f0000"]
    cmap1  = LinearSegmentedColormap.from_list("mycmap", color1)

    nx = len(XTicks)
    ny = len(YTicks)

    #plot
    im = ax.imshow(ResultList, aspect = "auto", cmap=cmap1)

    # 添加颜色条
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="2%", pad=0.1)  # 设置 colorbar 的宽度和位置
    cb = plt.colorbar(im, cax=cax)

    # 设置坐标轴
    ax.set_xticks(np.arange(0, nx, 1), XTicks)
    ax.set_yticks(np.arange(0, ny, 1), YTicks)
    ax.xaxis.set_major_locator(MaxNLocator(nbins=10))
    ax.yaxis.set_major_locator(MaxNLocator(nbins=10))
    
    # 设置其他标签和标题等
    ax.set_xlabel('distance [$\mu$m]', fontsize=10)
    ax.set_ylabel('radius [$\mu$m]', fontsize=10)
    ax.set_title('THzIntensity [a.u.]', fontsize=10)




def plotSpectrum(ax, XTicks, YTicks, ResultList):
    #set colormap
    color1 = ["#ffffff","#ff0000","#7f0000"]
    cmap1  = LinearSegmentedColormap.from_list("mycmap", color1)

    nx = len(XTicks)
    ny = len(YTicks)

    #plot
    im = ax.imshow(ResultList, aspect = "auto", cmap=cmap1, origin='lower')

    # 添加颜色条
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="2%", pad=0.1)  # 设置 colorbar 的宽度和位置
    cb = plt.colorbar(im, cax=cax)

    # 设置坐标轴
    ax.set_xticks(np.arange(0, nx, 1), XTicks)
    ax.set_yticks(np.arange(0, ny, 1), YTicks)
    ax.xaxis.set_major_locator(MaxNLocator(nbins=10))
    ax.yaxis.set_major_locator(MaxNLocator(nbins=10))
    
    # 设置其他标签和标题等
    ax.set_xlabel('distance [$\mu$m]', fontsize=10)
    ax.set_ylabel('frequency [$THz$]', fontsize=10)
    ax.set_title('THzSpectrum [a.u.]', fontsize=10)


def plotAxial(ax, XTicks, ResultList):
    #plot
    ax.plot(XTicks, ResultList)
    # 设置其他标签和标题等
    ax.set_xlabel('distance [$\mu$m]', fontsize=10)
    ax.set_ylabel('Amplitude[a.u.]', fontsize=10)
    ax.set_title("THzIntensity(r=0, z)", fontsize=10)


def plotFre(ax, XTicks, ResultList):
    #plot
    ax.plot(XTicks, ResultList)

    # 设置其他标签和标题等
    ax.set_xlabel('fre [$THz$]', fontsize=10)
    ax.set_ylabel('amp[a.u.]', fontsize=10)
    ax.set_title("THzSpectrum(r=0)", fontsize=10)